package com.bruce.starts.demo.factory;

import org.apache.commons.lang.StringUtils;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;

public class TreeFactory<T>{
    protected Function<T, String> keyFunc;
    protected Function<T, String> parentKeyFunc;
    protected BiConsumer<T, List<T>> childrenFunc;
    protected Comparator<T> comparator;


    public Function<T, String> getKeyFunc() {
        return keyFunc;
    }

    public Function<T, String> getParentKeyFunc() {
        return parentKeyFunc;
    }

    public BiConsumer<T, List<T>> getChildrenFunc() {
        return childrenFunc;
    }

    public Comparator<T> getComparator() {
        return comparator;
    }

    public TreeFactory<T> keyFunc(Function<T, String> getKeyFunc) {
        this.keyFunc = getKeyFunc;
        return this;
    }

    public TreeFactory<T> parentKeyFunc(Function<T, String> getParentKeyFunc) {
        this.parentKeyFunc = getParentKeyFunc;
        return this;
    }

    public TreeFactory<T> setChildrenFunc(BiConsumer<T, List<T>> setChildrenFunc) {
        this.childrenFunc = setChildrenFunc;
        return this;
    }

    public TreeFactory<T> comparator(Comparator<T> comparator) {
        this.comparator = comparator;
        return this;
    }

    @SuppressWarnings("unchecked")
    public List<T> create(Collection<? extends T> data) {
        if (comparator != null) {
            List<T> dataList = data instanceof List ? (List<T>) data : new ArrayList<>(data);
            dataList.sort(comparator);
            data = dataList;
        }
        Map<Object, TreeNode<T>> nodeMap = new HashMap<>();
        List<TreeNode<T>> nodes = new LinkedList<>();
        for (T e : data) {
            final String key = keyFunc.apply(e);
            TreeNode<T> node = nodeMap.computeIfAbsent(key, TreeNode::new);
            node.data = e;
            final String parentKey = parentKeyFunc.apply(e);
            if (StringUtils.isNotBlank(parentKey)) {
                TreeNode<T> parentNode = node.parent = nodeMap.computeIfAbsent(parentKey, TreeNode::new);
                parentNode.addChild(e);
            }
            nodes.add(node);
        }
        List<T> roots = new LinkedList<>();
        for (TreeNode<T> node : nodes) {
            node.setChildren(childrenFunc);
            T e = node.data;
            if (node.isRoot() && e != null) {
                roots.add(e);
            }
        }
        return roots;
    }

    public static class TreeNode<T> {
        protected final Object key;
        protected TreeNode<T> parent;
        protected T data;
        protected List<T> children;

        public TreeNode(Object key) {
            this.key = key;
        }

        protected void addChild(T child) {
            if (children == null) {
                children = new LinkedList<>();
            }
            children.add(child);
        }

        protected void setChildren(BiConsumer<T, List<T>> setChildrenFunc) {
            if (data != null) {
                setChildrenFunc.accept(data, children);
            }
        }

        public boolean isRoot() {
            return parent == null;
        }

        public Object getKey() {
            return key;
        }

        public TreeNode<T> getParent() {
            return parent;
        }

        public T getData() {
            return data;
        }

        public List<T> getChildren() {
            return children;
        }
    }
}
